clear; clc; close all;

% 读入原图（任意大小彩色或灰度）
I = imread('Lena.tif');
if size(I,3) == 3
    I = rgb2gray(I);
end

s = 256;

% 缩放到256x256
I = imresize(I, [s s]);

% 转double，方便后续计算
I = im2double(I);

% 调用加密函数
I_encrypted = encrypt_image(I);

% 保存加密后的图像
imwrite(I_encrypted, 'Lena_encrypted.tif');

% 显示结果
figure;
subplot(1,2,1);
imshow(I);
title('Original Image');

subplot(1,2,2);
imshow(I_encrypted);
title('Encrypted Image');



%% ———————— 加密函数及子函数 ————————
function I_encrypted = encrypt_image(I)
    % 输入：I - 原始256x256灰度图像，double格式，范围[0,1]
    % 输出：I_encrypted - 加密后图像，uint8格式，大小side_len x side_len

    % 1. DWT稀疏变换，取低频系数
    [LL,~,~,~] = dwt2(I,'haar');
    I1 = LL;
    I1_vec = I1(:);

    % 2. 设置压缩比例并计算压缩后长度M，保证为完全平方数
    compression_ratio = 0.8;
    vec_len = numel(I1_vec);
    M = round(compression_ratio * vec_len);
    M = floor(sqrt(M))^2;  % 保证M为完全平方数
    side_len = sqrt(M);

    % 3. 3D超混沌系统参数
    params.a = 1; params.b = 2*pi; params.c = 11.5;
    params.x0 = 0.4; params.y0 = 0.8; params.z0 = 0.6;

    % 4. 生成超混沌序列（前200步丢弃）
    [Xall, Yall, Zall] = generate_3D_SIMM_sequences(M+200, params);
    X = Xall(201:end); Y = Yall(201:end); Z = Zall(201:end);

    % 5. 构造Hadamard测量矩阵
    N_hadamard = 2^nextpow2(vec_len);
    H = hadamard(N_hadamard);
    H = H(:,1:vec_len);
    [~, idx] = sort(X, 'descend');
    Phi = H(idx(1:M), :);

    % 6. 压缩采样并量化
    I2_vec = Phi * I1_vec;
    I2q = uint8(255 * mat2gray(I2_vec));

    % 7. 改进Arnold置乱参数
    a_scramble = mod(floor(abs(X*1e6)), 256);
    b_scramble = mod(floor(abs(Y*1e6)), 256);

    % 8. 置乱
    scrambled_vec = improved_arnold_scramble(I2q, a_scramble(1), b_scramble(1));

    % 9. 转换为矩阵并转置
    scrambled_mat = reshape(scrambled_vec, side_len, side_len)';
    scrambled_vec = scrambled_mat(:);

    % 10. 像素扩散序列
    S1 = mod(floor(abs(Y*1e6)),256);
    S2 = mod(floor(abs(Z*1e6)),256);

    % 11. 像素扩散加密
    encrypted_vec = diffusion(scrambled_vec, S1, S2);

    % 12. 输出加密图像矩阵
    I_encrypted = reshape(encrypted_vec, side_len, side_len);
end


function [X,Y,Z] = generate_3D_SIMM_sequences(M, p)
    a = p.a; b = p.b; c = p.c;
    x0 = p.x0; y0 = p.y0; z0 = p.z0;
    X = zeros(1,M); Y = zeros(1,M); Z = zeros(1,M);
    X(1) = x0; Y(1) = y0; Z(1) = z0;
    for i=1:M-1
        X(i+1) = a*sin(b*Z(i))*sin(c/X(i));
        Y(i+1) = a*sin(b*X(i+1))*sin(c/Y(i));
        Z(i+1) = a*sin(b*Y(i+1))*sin(c/Z(i));
    end
end

function scrambled = improved_arnold_scramble(vec, a, b)
    N = numel(vec);
    scrambled = zeros(size(vec));
    for j=1:N
        q = mod(b + (a*b + 1)*j, N);
        idx = mod(q, N) + 1;
        scrambled(idx) = vec(j);
    end
    scrambled = uint8(scrambled);
end

function out_vec = diffusion(in_vec, S1, S2)
    N = length(in_vec);
    C = zeros(size(in_vec));
    C(1) = mod(double(in_vec(1)) + double(S1(1)), 256);
    for i=2:N
        C(i) = mod(C(i-1) + double(in_vec(i)) + double(S1(i)), 256);
    end
    C2 = zeros(size(C));
    C2(end) = C(end);
    for i=N-1:-1:1
        C2(i) = mod(C2(i+1) + double(C(i)) + double(S2(i)), 256);
    end
    out_vec = uint8(C2);
end
